'''
Created on May 17, 2011

@author: jlucasgl (jlucas.gl@gmail.com)

'''

from yacf.MiddleEnd.Loop.Analysis import ParametrizeLoopTool
from yacf.Backends.Common.Mutators.AbstractMutator import AbstractMutator

from yacf.Tools.Tree import ReplaceTool, getContainerAttribute
from yacf.Frontend.Shortcuts import getCurrentLanguageAst
from yacf.Backends.Common.Visitors.GenericVisitors import GenericFilterVisitor
c_ast = getCurrentLanguageAst()

from yacf.Frontend.Omp import omp_ast

class LoopTilingFilter(GenericFilterVisitor):
    """ Returns the first node with the given attribute
    """
    def __init__(self):
        def condition(node):
            if type(node) == omp_ast.llcTiling:
                return True
            return False
        super(LoopTilingFilter, self).__init__(condition_func = condition)

class LoopFilter(GenericFilterVisitor):
    """ Returns the first node with the given attribute
    """
    def __init__(self):
        def condition(node):
            if type(node) == c_ast.For:
                return True
            return False
        super(LoopFilter, self).__init__(condition_func = condition)

class IDFilter(GenericFilterVisitor):
    """ Returns the first node with the given attribute
    """
    def __init__(self):
        def condition(node):
            if type(node) == c_ast.ID:
                return True
            return False
        super(IDFilter, self).__init__(condition_func = condition)

class IDMutator(AbstractMutator):
    def __init__(self, old_id, new_id):
        self.old_id = old_id
        self.new_id = new_id

    def filter(self, ast):
        """ """
        raise NotImplemented

    def filter_iterator(self, ast):
        """ Fast filter """
        return NotImplemented

    def fast_filter(self, ast):
        """ Fast filter , looking for compound statement """
        return IDFilter().dfs_iter(ast)
    
    def mutatorFunction(self, id_found):
        if (id_found.name == self.old_id.name):
            # print " ID found in [", id_found.parent, "]"  
            setattr(id_found.parent, getContainerAttribute(id_found), self.new_id)


class SizeRequiredNotFound(Exception):
    def __init__(self):
        pass

    def __str__(self):
        return "The class constructor expects an argument of size"


class LoopTilingMutator(AbstractMutator):
    def __init__(self, size = None):
        self.size = size
        super(LoopTilingMutator, self).__init__()

    def filter(self, ast):
        """  """
        raise NotImplemented

    def filter_iterator(self, ast):
        """ Fast filter  """
        return NotImplemented

    def fast_filter(self, ast):
        """ Fast filter , looking for binary expressions """
        return LoopTilingFilter().dfs_iter(ast)
    
    def mutatorFunction(self, loopTiling_directive_ast):
    
        for_ast = LoopFilter().apply(loopTiling_directive_ast)
        assert isinstance(for_ast, c_ast.For)

        # Parametrize
        parametrizer = ParametrizeLoopTool()
        loop_parameters = parametrizer.apply(for_ast,)
        
        # Get llc node
        llc_node = loopTiling_directive_ast.parent

        #------------------------------------------------------------------
        # from yacf.Tools.Debug import DotDebugTool
        # DotDebugTool(highlight = [for_ast,]).apply(llc_node.parent.parent)
        #------------------------------------------------------------------

        # Get sizes of the tiling loop blocks
        try:
            self.size = loopTiling_directive_ast.size
        except AttributeError:
            if self.size is None:
                raise SizeRequiredNotFound()
        
        #=============================================================================
        # Modify parameters with the tiling
        #=============================================================================
        # 1. Build new for loop changing its next node
        #=============================================================================
        next_for = for_ast.next

        new_next_for = c_ast.Assignment(op = '+=',
                                       lvalue = loop_parameters['loop_variable'],
                                       rvalue = c_ast.Constant(type = 'int', value = str(self.size))
                                      )
        setattr(next_for.parent, getContainerAttribute(next_for), new_next_for)


        #=============================================================================
        # 2. Apply IDMutator to original for statement with to change loop_variable id
        #=============================================================================
        x_var_name = "__" + str(loop_parameters['loop_variable']) + "__"
        x_id_ast = c_ast.ID(name = x_var_name)

        IDMutator(loop_parameters['loop_variable'], x_id_ast).fast_apply_all(for_ast.stmt)

        
        #=============================================================================
        # 3. Build new internal for loop
        #=============================================================================
        #    3.1 Build the ternary operator parameters for last inner loop iteration
        #=============================================================================
        next_iteration = c_ast.BinaryOp(op = '+',
                                            left = loop_parameters['loop_variable'],
                                            right = c_ast.Constant(type = 'int', value = str(self.size))
                                    )

        last_iteration = loop_parameters['last_iteration']

        ternary_condition = c_ast.BinaryOp(op = '>',
                                            left = last_iteration,
                                            right = next_iteration
                                        )

        #=============================================================================
        #    3.2 Build the internal for loop with new loop_variable as index
        #=============================================================================
        cond_node = c_ast.BinaryOp(op = '<',
                                        left = x_id_ast,
                                        right = c_ast.TernaryOp(cond = ternary_condition, 
                                                                iftrue = next_iteration, 
                                                                iffalse = last_iteration
                                                            )
                                )

        init_node = c_ast.Assignment(op = '=',
                                        lvalue = x_id_ast,
                                        rvalue = loop_parameters['loop_variable']
                                    )

        #iterator = c_ast.UnaryOp(op='p++', expr = x_id_ast)

        internal_loop_parameters = {}
        internal_loop_parameters['init_node'] = init_node
        internal_loop_parameters['cond_node'] = cond_node
        internal_loop_parameters['loop_variable'] = x_id_ast
        internal_loop_parameters['stride'] = 1

        parametrizer_internal = ParametrizeLoopTool()
        parametrizer_internal.update(internal_loop_parameters)
        internal_for_ast = parametrizer_internal.buildLoop()

        #print internal_for_ast
        internal_for_ast.stmt = for_ast.stmt

        #DotDebugTool(highlight = []).apply(internal_for_ast)

        #=============================================================================
        #    3.3 Attach internal for loop to original for loop
        #=============================================================================
        new_internal_compound = c_ast.Compound(block_items = [internal_for_ast,])
        setattr(for_ast, getContainerAttribute(for_ast.stmt), new_internal_compound)

        #=============================================================================
        # 4. Replace and update the AST
        #=============================================================================
        #    4.1 Replace the forLoop with the tiling compound statement
        #=============================================================================
        ReplaceTool(new_node = for_ast, 
                    old_node = llc_node
                    ).apply(llc_node.parent, 
                            'block_items'
                            )

        #=============================================================================
        #    4.2 Finally, build declaration of new loop variable
        #=============================================================================
        x_variable_decl  =  c_ast.Decl(name = x_var_name,
                                        quals = [],
                                        storage = [],
                                        funcspec = [],
                                        type = c_ast.TypeDecl(declname = x_var_name, 
                                                        quals = [], 
                                                        type = c_ast.IdentifierType(names = ['int'])
                                                    ),
                                        init = None,
                                        bitsize = None
                                    )

        setattr(llc_node.parent, "block_items", [x_variable_decl, ] + llc_node.parent.block_items)
        #for_ast.parent.block_items = [x_variable_decl, ] + for_ast.parent.block_items

        #------------------------------------------------------------------
        # from yacf.Tools.Debug import DotDebugTool
        # DotDebugTool(highlight = [for_ast, x_variable_decl]).apply(llc_node.parent.parent)
        #------------------------------------------------------------------


